/* ************************************************************************* **
**    OpenVFIFE - Open System for Vector Form Instrinsic                     **
**                Finite Element Method (VFIFE)                              **
**                GinkGo(Tan Biao)                                           **
**                                                                           **
**                                                                           **
** (C) Copyright 2021, The GinkGo(Tan Biao). All Rights Reserved.            **
**                                                                           **
** Commercial use of this program without express permission of              **
** GinkGo(Tan Biao), is strictly prohibited.  See                            **
** file 'COPYRIGHT'  in main directory for information on usage and          **
** redistribution,  and for a DISCLAIMER OF ALL WARRANTIES.                  **
**                                                                           **
** Developed by:                                                             **
**      Tan Biao (ginkgoltd@outlook.com)                                     **
**                                                                           **
** ************************************************************************* */

// $Date: 2020-05-11 $
// Written: Tan Biao
// Revised:
//
// Purpose: This file contains the class definition for Particle

// The interface:
//

#include "../headers/particle.h"
using namespace std;


/*****************************************************************************/
/*                              Helper Functions                             */
/*****************************************************************************/
inline void replaceStdArray(StdArray6d &arr1, const StdArray6d &arr2)
{
    for (int i = 0; i < 6; i++)
    {
        arr1[i] = arr2[i];
    }
};

inline void addStdArray(StdArray6d &arr1, const StdArray6d &arr2)
{
    for (int i = 0; i < 6; i++)
    {
        arr1[i] += arr2[i];
    }
};


/*****************************************************************************/
/*                                  Particle                                 */
/*****************************************************************************/
inline void Particle::initiate()
{
    mass_.lumped_mass = {0, 0, 0, 0, 0, 0};
    mass_.translate_mass_matrix.setZero();
    mass_.rotation_mass_matrix.setZero();
    for (int i = 0; i < 6; i++)
    {
        force_[i] = 0;
        displace_[i] = 0;
        velocity_[i] = 0;
        accelerate_[i] = 0;
        previous_displace_[i] = 0;
        previous_velocity_[i] = 0;
        previous_accelerate_[i] = 0;
        current_position_[i] = coordinate_[i];
        previous_position_[i] = coordinate_[i];
        dof_index_[i] = false;
    }
    translate_mass_matrix_on_ = false;
    rotation_mass_matrix_on_ = false;
}

Particle::Particle(int id)
{
    id_ = id;
    initiate();
}

Particle::Particle(int id, double x, double y)
{
    id_ = id;
    coordinate_ = {0, 0, 0, 0, 0, 0};
    coordinate_[0] = x;
    coordinate_[1] = y;
    initiate();
}

Particle::Particle(int id, double x, double y, double z)
{
    id_ = id;
    coordinate_ = {0, 0, 0, 0, 0, 0};
    coordinate_[0] = x;
    coordinate_[1] = y;
    coordinate_[2] = z;
    initiate();
}

inline void Particle::checkMass()
{
    // lumped mass
    for (int i = 0; i < 6; i++)
    {
        if (dof_index_[i] && (mass_.lumped_mass[i] < 0))
        {
            std::cerr << "Warning: Negative lumped_mass!\n";
        }
    }
}

void Particle::info() const
{
    cout << "Particle ID: " << id_ << endl;
    cout << "coordinate: (";
    for (int i = 0; i < 6; i++)
    {
        cout << coordinate_[i] << ", ";
    }
    cout << ")" << endl;

    cout << "dofs: ";
    for (auto iter = dof_key_.begin(); iter != dof_key_.end(); ++iter)
    {
        cout << *iter << ", ";
    }
    cout << endl;

    cout << "mass:\n" << "lumped_mass: ";
    for (int i = 0; i < 6; i++)
    {
        cout << mass_.lumped_mass[i] << ", ";
    }
    cout << endl;
    cout << "consistent translate mass:" << endl;
    cout << mass_.translate_mass_matrix << endl;

    cout << "consistent rotation mass:" << endl;
    cout << mass_.rotation_mass_matrix << endl;

    cout << "force: ";
    for (int i = 0; i < 6; i++)
    {
        cout << force_[i] << ", ";
    }
    cout << "\n" << endl;
}

int Particle::id() const
{
    return id_;
}

const StdArray6d* Particle::coordinate() const
{
    return &coordinate_;
}
const Mass* Particle::mass() const
{
    return &mass_;
}

const StdArray6d* Particle::force() const
{
    return &force_;
}

const StdArray6d* Particle::displace() const
{
    return &displace_;
}


const StdArray6d* Particle::velocity() const
{
    return &velocity_;
}

const StdArray6d* Particle::accelerate() const
{
    return &accelerate_;
}


const StdArray6d* Particle::previous_displace() const
{
    return &previous_displace_;
}

const StdArray6d* Particle::previous_velocity() const
{
    return &previous_velocity_;
}

const StdArray6d* Particle::previous_accelerate() const
{
    return &previous_accelerate_;
}

const StdArray6d* Particle::current_position() const
{
    return &current_position_;
}

const StdArray6d* Particle::previous_position() const
{
    return &previous_position_;
}

const StdArray6b* Particle::dof_index() const
{
    return &dof_index_;
}

void Particle::setCoordinate(const StdArray6d &val)
{
    replaceStdArray(coordinate_, val);
}

void Particle::activateDof(string dof)
{
    if (dof == "Ux")
    {
        dof_key_.insert("Ux");
        dof_index_[0] = true;
    }
    else if (dof == "Uy")
    {
        dof_key_.insert("Uy");
        dof_index_[1] = true;
    }
    else if (dof == "Uz")
    {
        dof_key_.insert("Uz");
        dof_index_[2] = true;
    }
    else if (dof == "Rotx")
    {
        dof_key_.insert("Rotx");
        dof_index_[3] = true;
    }
    else if (dof == "Roty")
    {
        dof_key_.insert("Roty");
        dof_index_[4] = true;
    }
    else if (dof == "Rotz")
    {
        dof_key_.insert("Rotz");
        dof_index_[5] = true;
    }
    else
    {
        cerr << "unknown dof key: " << dof << endl;
    }
}

void Particle::deactivateDof(string dof)
{
    if (dof == "Ux")
    {
        dof_key_.erase("Ux");
        dof_index_[0] = false;
    }
    else if (dof == "Uy")
    {
        dof_key_.erase("Uy");
        dof_index_[1] = false;
    }
    else if (dof == "Uz")
    {
        dof_key_.erase("Uz");
        dof_index_[2] = false;
    }
    else if (dof == "Rotx")
    {
        dof_key_.erase("Rotx");
        dof_index_[3] = false;
    }
    else if (dof == "Roty")
    {
        dof_key_.erase("Roty");
        dof_index_[4] = false;
    }
    else if (dof == "Rotz")
    {
        dof_key_.erase("Rotz");
        dof_index_[5] = false;
    }
    else
    {
        cerr << "unknown dof key: " << dof << endl;
    }
}

void Particle::constraintDof(const DOF &dof)
{
    for (int i = 0; i < 6; i++)
    {
        if ((dof_index_[i]) && dof.key[i])
        {
            displace_[i] = dof.val[i];
            previous_displace_[i] = dof.val[i];
            velocity_[i] = 0;
            previous_velocity_[i] = 0;
            accelerate_[i] = 0;
            previous_accelerate_[i] = 0;
        }
    }
    updatePosition();
}

void Particle::setLumpedMass(const StdArray6d &val)
{
    replaceStdArray(mass_.lumped_mass, val);
    checkMass();
}

void Particle::addLumpedMass(const StdArray6d &val)
{
    addStdArray(mass_.lumped_mass, val);
    checkMass();
}

void Particle::clearLumpedMass()
{
    for (int i = 0; i < 6; i++)
    {
        mass_.lumped_mass[i] = 0;
    }
}

void Particle::setTranslateMass(const Eigen::Matrix3d &mat)
{
    translate_mass_matrix_on_ = true;
    mass_.translate_mass_matrix << mat;
    checkMass();
}

void Particle::addTranslateMass(const Eigen::Matrix3d &mat)
{
    translate_mass_matrix_on_ = true;
    mass_.translate_mass_matrix += mat;
    checkMass();
}

void Particle::clearTranslateMass()
{
    translate_mass_matrix_on_ = false;
    mass_.translate_mass_matrix.setZero();
}

void Particle::setInertiaMass(const Eigen::Matrix3d &mat)
{
    rotation_mass_matrix_on_ = true;
    mass_.rotation_mass_matrix << mat;
    checkMass();
}

void Particle::addInertiaMass(const Eigen::Matrix3d &mat)
{
    rotation_mass_matrix_on_ = true;
    mass_.rotation_mass_matrix += mat;
    checkMass();
}

void Particle::clearInertiaMass()
{
    rotation_mass_matrix_on_ = false;
    mass_.rotation_mass_matrix.setZero();
}

void Particle::setForce(const StdArray6d &val)
{
    replaceStdArray(force_, val);
}

void Particle::addForce(const StdArray6d &val)
{
    addStdArray(force_, val);
}

void Particle::clearForce()
{
    for (int i = 0; i < 6; i++)
    {
        force_[i] = 0;
    }
}

void Particle::setDisplace(const StdArray6d &val)
{
    replaceStdArray(displace_, val);
}

void Particle::addDisplace(const StdArray6d &val)
{
    addStdArray(displace_, val);
}

void Particle::clearDisplace()
{
    for (int i = 0; i < 6; i++)
    {
        displace_[i] = 0;
    }
}

void Particle::setVelocity(const StdArray6d &val)
{
    replaceStdArray(velocity_, val);
}

void Particle::addVelocity(const StdArray6d &val)
{
    addStdArray(velocity_, val);
}

void Particle::clearVelocity()
{
    for (int i = 0; i < 6; i++)
    {
        velocity_[i] = 0;
    }
}

void Particle::setAccelerate(const StdArray6d &val)
{
    replaceStdArray(accelerate_, val);
}

void Particle::addAccelerate(const StdArray6d &val)
{
    addStdArray(accelerate_, val);
}

void Particle::clearAccelerate()
{
    for (int i = 0; i < 6; i++)
    {
        accelerate_[i] = 0;
    }
}

void Particle::setPreviousDisplace(const StdArray6d &val)
{
    replaceStdArray(previous_displace_, val);
}

void Particle::addPreviousDisplace(const StdArray6d &val)
{
    addStdArray(previous_displace_, val);
}

void Particle::clearPreviousDisplace()
{
    for (int i = 0; i < 6; i++)
    {
        previous_displace_[i] = 0;
    }
}

void Particle::setPreviousVelocity(const StdArray6d &val)
{
    replaceStdArray(previous_velocity_, val);
}

void Particle::addPreviousVelocity(const StdArray6d &val)
{
    addStdArray(previous_velocity_, val);
}

void Particle::clearPreviousVelocity()
{
    for (int i = 0; i < 6; i++)
    {
        if (dof_index_[i])
        {
            previous_velocity_[i] = 0;
        }
    }
}

void Particle::setPreviousAccelerate(const StdArray6d &val)
{
    replaceStdArray(previous_accelerate_, val);
}

void Particle::addPreviousAccelerate(const StdArray6d &val)
{
    addStdArray(previous_accelerate_, val);
}

void Particle::clearPreviousAccelerate()
{
    for (int i = 0; i < 6; i++)
    {
        previous_accelerate_[i] = 0;
    }
}

void Particle::updatePosition()
{
    for (int i = 0; i < 6; i++)
    {
        current_position_[i] = coordinate_[i] + displace_[i];
        previous_position_[i] = coordinate_[i] + previous_displace_[i];
    }
}

void Particle::clearAll()
{
    clearDisplace();
    clearVelocity();
    clearAccelerate();
    clearPreviousDisplace();
    clearPreviousVelocity();
    clearPreviousAccelerate();
    clearForce();
    // reset position
    for (int i = 0; i < 6; i++)
    {
        current_position_[i] = coordinate_[i];
        previous_position_[i] = coordinate_[i];
    }
}

void Particle::solveFirstStep(double h, double zeta)
{
    if ((!translate_mass_matrix_on_) && (!rotation_mass_matrix_on_))
    {
        for (int i = 0; i < 6; i++)
        {
            if (dof_index_[i])
            {
                accelerate_[i] = force_[i] / mass_.lumped_mass[i];
                previous_displace_[i] = accelerate_[i] * h * h / 2.0 -
                                    h * (1 + 0.5 * zeta * h) * velocity_[i] +
                                    displace_[i];
            }
        }
    }
    else if (rotation_mass_matrix_on_)
    {
        for (int i = 0; i < 3; i++)
        {
            if (dof_index_[i])
            {
                accelerate_[i] = force_[i] / mass_.lumped_mass[i];
                previous_displace_[i] = accelerate_[i] * h * h / 2.0 -
                                    h * (1 + 0.5 * zeta * h) * velocity_[i] +
                                    displace_[i];
            }
        }

        Eigen::Vector3d moment, tmp;
        moment << force_[3], force_[4], force_[5];
        tmp = mass_.rotation_mass_matrix.inverse() * moment;

        for (int i = 3; i < 6; i++)
        {
            accelerate_[i] = tmp(i-3);
            previous_displace_[i] = accelerate_[i] * h * h / 2.0 -
                                h * (1 + 0.5 * zeta * h) * velocity_[i] +
                                displace_[i];

        }
    }

    updatePosition();
}

void Particle::solve(double h, double zeta)
{
    const double c1 = 1.0 + zeta * h / 2.0;
    const double c2 = 1.0 - zeta * h / 2.0;
    const double h2 = h * h;

    StdArray6d disp = {};
    StdArray6d velo = {};
    StdArray6d acc = {};

    if ((!translate_mass_matrix_on_) && (!rotation_mass_matrix_on_))
    {
        for (int i = 0; i < 6; i++)
        {
            if (dof_index_[i])
            {
                disp[i] = (2.0 * displace_[i] - c2 * previous_displace_[i] +
                          h2 * force_[i] / mass_.lumped_mass[i]) / c1;
                velo[i] = (disp[i] - previous_displace_[i]) / h / 2.0;
                acc[i] = (disp[i] - 2.0 * displace_[i] +
                          previous_displace_[i]) / h2;
            }
        }
    }
    else if (rotation_mass_matrix_on_)
    {
        for (int i = 0; i < 3; i++)
        {
            if (dof_index_[i])
            {
                disp[i] = (2.0 * displace_[i] - c2 * previous_displace_[i] +
                          h2 * force_[i] / mass_.lumped_mass[i]) / c1;
                velo[i] = (disp[i] - previous_displace_[i]) / h / 2.0;
                acc[i] = (disp[i] - 2.0 * displace_[i] +
                          previous_displace_[i]) / h2;
            }
        }

        Eigen::Vector3d moment, tmp;
        moment << force_[3], force_[4], force_[5];
        tmp = mass_.rotation_mass_matrix.inverse() * moment;

        for (int i = 3; i < 6; i++)
        {
            disp[i] = (2.0 * displace_[i] - c2 * previous_displace_[i] +
                       h2 * tmp(i-3)) / c1;
            velo[i] = (disp[i] - previous_displace_[i]) / h / 2.0;

            acc[i] = (disp[i] - 2.0 * displace_[i] +
                      previous_displace_[i]) / h2;
        }
    }

    setPreviousDisplace(displace_);
    setPreviousVelocity(velocity_);
    setPreviousAccelerate(accelerate_);
    setDisplace(disp);
    setVelocity(velo);
    setAccelerate(acc);
    updatePosition();
}

void Particle::outputForce(fstream &fout) const
{
    fout << id_ << ", ";
    for (int i = 0; i < 5; i++)
    {
        fout << force_[i] << ", ";
    }
    fout << force_[5] << endl;
}

void Particle::outputReactionForce(fstream &fout) const
{
    fout << id_ << ", ";
    for (int i = 0; i < 5; i++)
    {
        fout << -force_[i] << ", ";
    }
    fout << -force_[5] << endl;
}

void Particle::outputDisplace(fstream &fout) const
{
    fout << id_ << ", ";
    for (int i = 0; i < 5; i++)
    {
        fout << displace_[i] << ", ";
    }
    fout << displace_[5] << endl;
}

void Particle::outputVelocity(fstream &fout) const
{
    fout << id_ << ", ";
    for (int i = 0; i < 5; i++)
    {
        fout << velocity_[i] << ", ";
    }
    fout << velocity_[5] << endl;
}

void Particle::outputAccelarate(fstream &fout) const
{
    fout << id_ << ", ";
    for (int i = 0; i < 5; i++)
    {
        fout << accelerate_[i] << ", ";
    }
    fout << accelerate_[5] << endl;
}
